#!/usr/bin/python
import cv2
import cv2.cv as cv
import numpy as np
import roslib; roslib.load_manifest('ardrone_joy_control')
import rospy
from sensor_msgs.msg import Image
import roslib.packages
from cv_bridge import CvBridge, CvBridgeError
from ardrone_autonomy.msg import Navdata
import os


class artificial_horizont (object):
    def __init__ (self,image_path):
        self.error=0
        #self.image_path= os.path.dirname(__file__)
        #print os.path.pardir
        #print self.image_path
        
        self.fixed_border = cv2.imread(image_path + "fijo.png")
        if self.fixed_border == None:
            rospy.loginfo(rospy.get_name()+ "Error al cargar " + image_path + "fijo.png")
            self.error=1
            
        self.fixed_border_mask = cv2.imread(image_path +"fijo-mask.png")
        if self.fixed_border_mask == None:
            rospy.loginfo(rospy.get_name()+ "Error al cargar " + image_path + "fijo-mask.png")
            self.error=1
        else:
            self.fixed_border_mask/=255
            
        self.banking = cv2.imread(image_path +"banqueo.png")
        if self.banking == None:
            rospy.loginfo(rospy.get_name()+ "Error al cargar " + image_path + "fijo.png")
            self.error=1

        self.banking_mask = cv2.imread(image_path +"banqueo-mask.png")
        if self.banking_mask == None:
            rospy.loginfo(rospy.get_name()+ "Error al cargar " + image_path + "banqueo-mask.png")
            self.error=1
        else:
            self.banking_mask/=255
            
        self.horizont = cv2.imread(image_path +"mundo.png")
        if self.horizont == None:
            rospy.loginfo(rospy.get_name()+ "Error al cargar " + image_path + "mundo.png")
            self.error=1
            
        self.horizont_mask = cv2.imread(image_path +"visible-mask.png")
        if self.horizont_mask == None:
            rospy.loginfo(rospy.get_name()+ "Error al cargar " + image_path + "visible-mask.png")
            self.error=1
        else:
            self.horizont_mask/=255
            
        if not self.error:
            self.fixed_border_h,    self.fixed_border_w     =   self.fixed_border.shape[:2]
            self.horizont_mask_h,   self.horizont_mask_w    =   self.horizont_mask.shape[:2]
            self.banking_h,         self.banking_w          =   self.banking.shape[:2]
            self.horizont_h,        self.horizont_w         =   self.horizont.shape[:2]
            
            self.horizont_full_mask =   np.zeros(self.fixed_border.shape,np.uint8)
            
            self.horizont_mask_x1=(self.fixed_border_w - self.horizont_mask_w)/2
            self.horizont_mask_x2=(self.fixed_border_w + self.horizont_mask_w)/2
            self.horizont_mask_y1=(self.fixed_border_h - self.horizont_mask_h)/2
            self.horizont_mask_y2=(self.fixed_border_h + self.horizont_mask_h)/2
            
            self.horizont_full_mask[self.horizont_mask_y1:self.horizont_mask_y2, self.horizont_mask_x1:self.horizont_mask_x2]=(1-self.horizont_mask)
            
            self.horizont_x1=(self.horizont_w - self.fixed_border_w)/2
            self.horizont_x2=(self.horizont_w + self.fixed_border_w)/2
            self.horizont_y1=(self.horizont_h - self.fixed_border_h)/2
            self.horizont_y2=(self.horizont_h + self.fixed_border_h)/2
            
            self.banking_visible=self.banking * self.banking_mask
            
            self.banking_mask_x1=(self.fixed_border_w - self.banking_w)/2
            self.banking_mask_x2=(self.fixed_border_w + self.banking_w)/2
            self.banking_mask_y1=(self.fixed_border_h - self.banking_h)/2
            self.banking_mask_y2=(self.fixed_border_h + self.banking_h)/2
            
            self.instrument=np.zeros(self.fixed_border.shape,np.uint8)

    def draw(self,pitch,roll):
        if not self.error:
            self.horizont_visible=self.horizont[self.horizont_y1-2*pitch:self.horizont_y2-2*pitch,self.horizont_x1:self.horizont_x2]
            
            
            self.horizont_visible[self.banking_mask_y1:self.banking_mask_y2,self.banking_mask_x1:self.banking_mask_x2]*=(1-self.banking_mask)
            
            self.horizont_visible[self.banking_mask_y1:self.banking_mask_y2,self.banking_mask_x1:self.banking_mask_x2]+=self.banking_visible
            self.rotation=cv2.getRotationMatrix2D((self.fixed_border_h/2,self.fixed_border_w/2), roll, 1)
            
            self.horizont_visible_rotated=cv2.warpAffine(self.horizont_visible,self.rotation,self.fixed_border.shape[:2])*(self.horizont_full_mask)*self.fixed_border_mask
            self.fixed_border_visible=self.fixed_border*(1-self.fixed_border_mask)
            
            self.instrument=self.horizont_visible_rotated+self.fixed_border_visible
        else:
            rospy.loginfo(rospy.get_name()+ "Hubo un error y no se puede crear el instrumento")
            self.instrument=np.zeros([10,10],np.uint8)
            
            
class node (object):
    def __init__ (self):
        self.bridge = CvBridge()
        self.image_path= roslib.packages.get_pkg_dir('ardrone_joy_control') + "/images/"
        self.hsi=artificial_horizont(self.image_path)
        rospy.init_node('ardrone_hsi_node', anonymous=True)    
        
        self.image_pub=rospy.Publisher('/ardrone_joy_control/hsi', Image)
        rospy.Subscriber("/ardrone/navdata", Navdata, self.process_navdata)
        rospy.spin()

    def process_navdata(self,data):
        self.hsi.draw(-data.rotY,data.rotX)

        source=self.hsi.instrument
        bitmap = cv.CreateImageHeader((source.shape[1], source.shape[0]), cv.IPL_DEPTH_8U, 3)
        cv.SetData(bitmap, source.tostring(),source.dtype.itemsize * 3 * source.shape[1])
        try:
            self.image_pub.publish(self.bridge.cv_to_imgmsg(bitmap, "bgr8"))
        except CvBridgeError, e:
            rospy.loginfo(rospy.get_name()+ e)
            

if __name__ == '__main__':
    my_node=node()
